package com.kly.aop;

import cn.hutool.core.util.StrUtil;
import com.google.common.util.concurrent.RateLimiter;
import com.kly.annotation.GuavaLimitAnno;
import com.kly.exception.ErrorCode;
import com.kly.exception.ThrowUtils;
import com.kly.manager.auth.StpKit;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @author admin
 * @date 2025/5/5
 */
@Aspect
@Component
@Slf4j
public class GuavaLimit {

    private static final Map<String, RateLimiter> map = new ConcurrentHashMap<>();

    @Pointcut("execution(* com.kly.controller.*.*(..))")
    private void pc() {
    }

    @Before("pc() && @annotation(guavaLimitAnno)")
    public void guavaLimit(GuavaLimitAnno guavaLimitAnno) {
        // 获取用户id
        long userId = StpKit.SPACE.getLoginIdAsLong();
        String key = guavaLimitAnno.key() + ":" + userId;
        double permitsPerSecond = guavaLimitAnno.permitsPerSecond();
        long timeout = guavaLimitAnno.timeout();
        TimeUnit timeunit = guavaLimitAnno.timeunit();
        String msg = guavaLimitAnno.msg();

        if (StrUtil.isBlank(key)) {
            key = "default";
        }
        RateLimiter rateLimiter = null;
    //     创建限流器
        if (!map.containsKey(key)) {
            rateLimiter = rateLimiter.create(permitsPerSecond);
            map.put(key, rateLimiter);
        }
        rateLimiter = map.get(key);

        boolean acquire = rateLimiter.tryAcquire(timeout, timeunit);
        ThrowUtils.throwIf(!acquire, ErrorCode.OPERATION_ERROR, msg);
    }
}
